from transformers import pipeline
import os

try:
    # 尝试在线加载模型
    ner_pipeline = pipeline(
        "token-classification",
        model="microsoft/ple-bert-base",
        timeout=60  # 增加超时时间
    )
except OSError:
    # 离线模式加载本地模型
    local_model_path = "./ple-bert-base"
    if os.path.exists(local_model_path):
        ner_pipeline = pipeline(
            "token-classification",
            model=local_model_path
        )
    else:
        raise RuntimeError(
            "无法连接HuggingFace且本地模型不存在\n"
            "请执行以下步骤：\n"
            "1. 检查网络连接\n"
            "2. 或运行以下命令下载模型：\n"
            "   huggingface-cli download microsoft/ple-bert-base --local-dir ./ple-bert-base"
        )


def apply_watermark(text, entity):
    """添加不可见水印的示例实现"""
    original = text[entity['start']:entity['end']]
    watermarked = f"{original}‹⁠⁠⁠ₓ⁠⁠›"  # 零宽空格字符
    return text[:entity['start']] + watermarked + text[entity['end']:]


def smart_redact(text):
    entities = ner_pipeline(text)

    # 按倒序处理避免偏移问题
    for entity in sorted(entities, key=lambda x: x['start'], reverse=True):
        if entity['entity_group'] in ['PHONE', 'ID_NUMBER']:
            replacement = '[REDACTED]'
        elif entity['entity_group'] == 'INTERNAL_CODE':
            replacement = apply_watermark(text, entity)
        else:
            continue

        text = text[:entity['start']] + replacement + text[entity['end']:]

    return text


# 测试（确保测试文本包含所有实体类型）
if __name__ == "__main__":
    test_text = "联系电话：138-1234-5678，身份证号：11010119900307765X"
    print(smart_redact(test_text))